import cv2
import numpy as np

# Read input video
cap = cv2.VideoCapture('test.mp4')

# Get frame count
n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

# Get width and height of video stream
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# Define the codec for output video
# fourcc = cv2.VideoWriter_fourcc(*'MJPG')
#
# # Set up output video
# out = cv2.VideoWriter('video_out.mp4', fourcc, fps, (w, h))

# Read first frame
_, prev = cap.read()

# Convert frame to grayscale
prev_gray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
# Pre-define transformation-store array
transforms = np.zeros((n_frames - 1, 3), np.float32)

for i in range(n_frames - 2):
    # Detect feature points in previous frame
    prev_pts = cv2.goodFeaturesToTrack(prev_gray,
                                       maxCorners=200,
                                       qualityLevel=0.01,
                                       minDistance=30,
                                       blockSize=3)

    # Read next frame
    success, curr = cap.read()
    if not success:
        break

        # Convert to grayscale
    curr_gray = cv2.cvtColor(curr, cv2.COLOR_BGR2GRAY)

    # Calculate optical flow (i.e. track feature points)
    curr_pts, status, err = cv2.calcOpticalFlowPyrLK(prev_gray, curr_gray, prev_pts, None)

    # Sanity check
    assert prev_pts.shape == curr_pts.shape

    # Filter only valid points
    idx = np.where(status == 1)[0]
    prev_pts = prev_pts[idx]
    curr_pts = curr_pts[idx]
    moveing = curr_pts - prev_pts
    x=moveing[...,0]
    y=moveing[...,1]
    print (np.mean(x),np.mean(y))

    # Find transformation matrix
    # m = cv2.estimateAffinePartial2D(prev_pts, curr_pts,True)[0] # will only work with OpenCV-3 or less
    # # Extract traslation
    # dx = m[0, 2]
    # dy = m[1, 2]
    #
    # # Extract rotation angle
    # da = np.arctan2(m[1, 0], m[0, 0])
    #
    # # Store transformation
    # transforms[i] = [dx, dy, da]

    # Move to next frame
    prev_gray = curr_gray
